import { MongoClient } from "mongodb";
import chalk from "chalk";
import {
    MONGO_URI,
    MONGO_DB,
    SOURCE_COLLECTION,
    SOURCE_TEMP,
    DOCUMENTS_SP_RAG,
} from "./config.js";

// 处理开头包含 @
const processSpcialChar = (str) => {
    if (str.indexOf("@") === 0) {
        return str.substr(1);
    }
    return str;
};
const removeBrackets = (str) => {
    return str.replace(/\【.*?\】/g, "").trim();
};

const removeText = (content, text) => {
    const regex = new RegExp(text, "g");
    return content.replace(regex, "") || content;
};
const removeSubstring = (str, start, end) => {
    let realEnd = end
    if (str[end] === '\n') {
        realEnd++
        console.log('是换行符');
    }
    return str.slice(0, start) + str.slice(realEnd);
};

async function updateData() {
    const client = new MongoClient(MONGO_URI, {
        useNewUrlParser: true,
        useUnifiedTopology: true,
    });

    try {
        // 连接到 MongoDB
        await client.connect();
        console.log(chalk.bold.yellow("MongoDB connection connected"));
        const db = client.db(MONGO_DB);
        const sourceCollection = db.collection(SOURCE_TEMP);
        const destinationCollection = db.collection(DOCUMENTS_SP_RAG);

        // 创建一个游标来遍历集合
        // title:{'$regex': '.*方案.*篇.*'}  content:{'$regex': '.*附件'}
        // 文章不规范：存在"附件:" 、"附件："两种情况

        // 查询annotions
        // {
        //     annotations: { $size: 1 },
        // }

        // 处理多篇的
        // {id: 18871}

        const cursor = sourceCollection.find(
            {
                annotations: { $size: 1 },
            },
            { noCursorTimeout: true }
        );

        for await (const doc of cursor) {
            const id = doc.data.id;
            // const exists = await destinationCollection.findOne({ id });
            // 检查是否已存在
            //   if (exists) {
            //     console.log(`fragment with id=${id} already exists`);
            //     continue;
            //   }
            const data = doc.data;
            let isCanDroped = false;
            let isCuted = false

            // 这里是去除@
            //   const title = processSpcialChar(data.title);
            //   const content = processSpcialChar(data.content);

            // 这里是去除中括号
            // const title = removeBrackets(data.title);
            // const content = removeBrackets(data.content);

            const content = data.content;
            const { annotations } = doc;
            let cleanedContent = content;
            // 虽然annotations 是数组 但是 数组长度大于1 仅为4个 且全部为 优质入库

            const result = annotations[0].result;

            const len = result.length;

            for (let i = 0; i < len; i++) {
                const item = result[i];
                // 原来是 7187个
                if (item.value?.choices && item.value?.choices.includes("低质舍弃")) {
                    isCanDroped = true;

                }
                if (
                    item.from_name === "highlight_labels" &&
                    item.value.start &&
                    item.value.end
                ) {
                    cleanedContent = removeSubstring(
                        cleanedContent,
                        item.value.start,
                        item.value.end
                    );
                    isCuted = true
                }
            }
            const cleanedDoc = {
                ...data,
                isCanDroped,
                isCuted,
                content: cleanedContent,
            };

            console.log('当前 isCanDroped: ', isCanDroped);


            // 更新到 SOURCE_COLLECTION
            await destinationCollection.insertOne({
                ...cleanedDoc,
            });
            console.log(chalk.green(`insertOne document: ${doc.id}`));
        }

        console.log(chalk.bold.yellow("Success: all docs inserted"));
    } catch (err) {
        console.error(err);
    } finally {
        console.log(chalk.bold.yellow("MongoDB connection closed"));
        await client.close();
    }
}

// 执行数据更新
updateData().catch(console.error);
